{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "990f673b",
   "metadata": {},
   "source": [
    "# HW4: Recipe RAG — BM25, Synthetic Queries, Phoenix\n",
    "\n",
    "\n",
    "<center>\n",
    "    <p style=\"text-align:left\">\n",
    "        <img alt=\"phoenix logo\" src=\"https://repository-images.githubusercontent.com/564072810/f3666cdf-cb3e-4056-8a25-27cb3e6b5848\" width=\"600\"/>\n",
    "        <br>\n",
    "        <a href=\"https://arize.com/docs/phoenix/\">Docs</a>\n",
    "        |\n",
    "        <a href=\"https://github.com/Arize-ai/phoenix\">GitHub</a>\n",
    "        |\n",
    "        <a href=\"https://arize-ai.slack.com/join/shared_invite/zt-2w57bhem8-hq24MB6u7yE_ZF_ilOYSBw#/shared-invite/email\">Community</a>\n",
    "    </p>\n",
    "</center>\n",
    "\n",
    "## 🎯 Assignment Overview\n",
    "In this assignment, you’ll build and evaluate a BM25-based RAG system for recipes, generate synthetic user queries with llm_generate, and trace everything in Phoenix for visibility and analysis.\n",
    "\n",
    "#### Workflow\n",
    "- 🧰 Setup — Install deps, set env vars, and import libs.\n",
    "\n",
    "- 📥 Load & structure data — Read RAW_recipes.csv, normalize fields (title, ingredients, steps, raw text).\n",
    "\n",
    "- 🔤 Tokenize & index — Use a number-preserving tokenizer (keeps 375 F, 9x13) and build/load a BM25 index.\n",
    "\n",
    "- 🧪 Generate synthetic queries — Use llm_generate to create realistic, technical cooking questions (+ validation & dedupe).\n",
    "\n",
    "- ☁️ Upload dataset to Phoenix — Map to input (query) and output (salient_fact) and upload as a Phoenix Dataset.\n",
    "\n",
    "- 🔎 Run retriever with tracing — For each query, call retrieve_bm25() (decorated) to emit rag.query → retrieval.bm25 spans.\n",
    "\n",
    "- 🧪 Experiments — Run a Phoenix Experiment (Recall@1/3/5, MRR) and compare runs.\n",
    "\n",
    "- 💾 Save results — Persist per-query outputs & summary to JSON for offline review.\n",
    "\n",
    "- 📊 Inspect in Phoenix — View spans, top-k documents, evaluation chips, and experiment summaries.\n",
    "\n",
    "#### Core Task: “Can the retriever surface the correct recipe?”\n",
    "You’ll evaluate whether the ground-truth recipe (the one that the synthetic question depends on) appears in the top-k results—especially top-1/3/5."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98c009a9",
   "metadata": {},
   "source": [
    "## Environment & Imports\n",
    "_Install deps (if needed), configure environment variables, and import libraries. Keep this section idempotent so re-running the notebook is safe._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "852cf7a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -qqqq 'arize-phoenix==11.21.0' rank-bm25 tqdm litellm python-dotenv kagglehub openinference-instrumentation-litellm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "379551e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ast\n",
    "import functools\n",
    "import json\n",
    "import math\n",
    "import os\n",
    "import pickle\n",
    "import random\n",
    "import re\n",
    "import string\n",
    "import time\n",
    "from contextlib import contextmanager, nullcontext\n",
    "from getpass import getpass\n",
    "from pathlib import Path\n",
    "from time import perf_counter\n",
    "from typing import Any, Dict, List, Optional, Union\n",
    "\n",
    "import kagglehub\n",
    "import nest_asyncio\n",
    "import pandas as pd\n",
    "from kagglehub import KaggleDatasetAdapter\n",
    "from openinference.semconv.trace import OpenInferenceSpanKindValues, SpanAttributes\n",
    "from rank_bm25 import BM25Okapi\n",
    "from tqdm import tqdm\n",
    "\n",
    "import phoenix as px\n",
    "from phoenix.client import AsyncClient\n",
    "from phoenix.client.experiments import run_experiment\n",
    "from phoenix.evals import OpenAIModel, RelevanceEvaluator, llm_generate, run_evals\n",
    "from phoenix.otel import register\n",
    "from phoenix.session.evaluation import get_retrieved_documents\n",
    "from phoenix.trace import DocumentEvaluations\n",
    "\n",
    "file_path = \"hw4/RAW_recipes.csv\"\n",
    "KAGGLE_DATASET = \"shuyangli94/food-com-recipes-and-user-interactions\"\n",
    "KAGGLE_FILE = \"RAW_recipes.csv\"\n",
    "LOCAL_RAW_CSV = \"hw4/data/RAW_recipes.csv\"\n",
    "OUTPUT_JSON = \"hw4/data/processed_recipes.json\"\n",
    "DEFAULT_INDEX_PATH = \"hw4/data/bm25_index.pkl\"\n",
    "PHOENIX_API_KEY = os.getenv(\"PHOENIX_API_KEY\")\n",
    "PHOENIX_ENDPOINT = os.getenv(\"PHOENIX_COLLECTOR_ENDPOINT\")\n",
    "tracer = None\n",
    "TEXT_FIELD = \"raw_text\"\n",
    "TITLE_FIELD = \"title\"\n",
    "ID_FIELD = \"id\"\n",
    "TOP_N = 200\n",
    "\n",
    "QUERY_MODEL = \"gpt-3.5-turbo-0125\"\n",
    "NUM_QUERIES = 100\n",
    "BATCH_SIZE = 3\n",
    "MAX_WORKERS = 5\n",
    "SEED = 42\n",
    "TIMEOUT_S = 60\n",
    "SYN_QUERY_PATH = \"hw4/data/synthetic_queries.json\"\n",
    "EVAL_RESULTS_PATH = \"hw4/results/retrieval_evaluation.json\"\n",
    "\n",
    "random.seed(SEED)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e7bf2bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "_JSON_RE = re.compile(r\"\\{.*\\}\", re.S)\n",
    "_WORD_RE = re.compile(r\"[a-zA-Z]+\")\n",
    "_TECH = re.compile(\n",
    "    r\"\\b\\d[\\d/\\.\\s]*(?:sec|second|min|minute|hour|hr|°[CF]|deg(?:rees)?\\s*[CF]|tsp|tbsp|cup|cups|g|gram|ml|°)\\b\",\n",
    "    re.I,\n",
    ")\n",
    "_TOKEN_RE = re.compile(\n",
    "    r\"\\d+\\s*[x×]\\s*\\d+\"\n",
    "    r\"|(?:\\d+/?\\d+)\"\n",
    "    r\"|(?:\\d+(?:\\.\\d+)?)\"\n",
    "    r\"|(?:°[fc])\"\n",
    "    r\"|[a-z]+\"\n",
    ")\n",
    "\n",
    "\n",
    "def tokenize(text: str) -> list[str]:\n",
    "    s = (text or \"\").lower()\n",
    "    s = s.replace(\"degrees f\", \"°f\").replace(\"degree f\", \"°f\")\n",
    "    s = s.replace(\"degrees c\", \"°c\").replace(\"degree c\", \"°c\")\n",
    "    s = s.replace(\"mins\", \"min\").replace(\"minutes\", \"min\")\n",
    "    s = s.replace(\"hours\", \"hr\").replace(\"hrs\", \"hr\")\n",
    "    return _TOKEN_RE.findall(s)\n",
    "\n",
    "\n",
    "def _to_text(x) -> str:\n",
    "    if isinstance(x, str):\n",
    "        return x\n",
    "    if x is None:\n",
    "        return \"\"\n",
    "    if isinstance(x, float) and math.isnan(x):\n",
    "        return \"\"\n",
    "    return str(x)\n",
    "\n",
    "\n",
    "def _jsonish(v):\n",
    "    if isinstance(v, str):\n",
    "        t = v.strip()\n",
    "        if t and (t[0] + t[-1] in (\"{}\", \"[]\")):\n",
    "            try:\n",
    "                return json.loads(t)\n",
    "            except Exception:\n",
    "                pass\n",
    "    return v\n",
    "\n",
    "\n",
    "def _norm_id(x):\n",
    "    return \"\" if x is None else str(x)\n",
    "\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5f6e372",
   "metadata": {},
   "source": [
    "## Step 0: Setup Phoenix"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff64a6da",
   "metadata": {},
   "source": [
    "First, let's set up Phoenix on our local machine. You should run these commands within your terminal in your chosen environment.\n",
    "\n",
    "(If you have already done this in a previous HW assignment, you are good to go.)\n",
    "\n",
    "**Boot up Phoenix on localhost**\n",
    "\n",
    "```phoenix serve```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f909c1e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def register_phoenix():\n",
    "    global tracer\n",
    "    if tracer:\n",
    "        return tracer\n",
    "    tracer_provider = register(project_name=\"hw4-rag\")\n",
    "    tracer = tracer_provider.get_tracer(__name__)\n",
    "    return tracer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f709081",
   "metadata": {},
   "outputs": [],
   "source": [
    "tracer = register_phoenix()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de1e31fd",
   "metadata": {},
   "source": [
    "#### Set OpenAI API Key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "555c3811",
   "metadata": {},
   "outputs": [],
   "source": [
    "if \"OPENAI_API_KEY\" not in os.environ:\n",
    "    os.environ[\"OPENAI_API_KEY\"] = getpass(\"🔑 Enter your OpenAI API key: \")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "606c3e5a",
   "metadata": {},
   "source": [
    "# Part 1: Create Your Retrieval Evaluation Dataset\n",
    "\n",
    "## Step1: Data Loading & Structuring\n",
    "_Load the recipe dataset, perform minimal cleaning, and structure the fields we care about (title, ingredients, steps, text). Avoid heavy processing here so experimentation is quick._\n",
    "- Load and clean the provided RAW_recipes.csv (~5,000 recipes)\n",
    "- Structure recipe data (ingredients, steps, tags, nutrition)\n",
    "- Select the ~200 longest recipes by text content for richer evaluation\n",
    "- Save as data/processed_recipes.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "482fa7f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_raw_recipes() -> pd.DataFrame:\n",
    "    \"\"\"\n",
    "    Try local CSV first. If missing, pull via kagglehub and also cache locally.\n",
    "    \"\"\"\n",
    "    if os.path.exists(LOCAL_RAW_CSV):\n",
    "        return pd.read_csv(LOCAL_RAW_CSV)\n",
    "    try:\n",
    "        df = kagglehub.load_dataset(\n",
    "            KaggleDatasetAdapter.PANDAS,\n",
    "            KAGGLE_DATASET,\n",
    "            KAGGLE_FILE,\n",
    "        )\n",
    "        Path(LOCAL_RAW_CSV).parent.mkdir(parents=True, exist_ok=True)\n",
    "        df.to_csv(LOCAL_RAW_CSV, index=False)\n",
    "        print(f\"✅ Saved raw CSV to {LOCAL_RAW_CSV}\")\n",
    "        return df\n",
    "    except ModuleNotFoundError as e:\n",
    "        raise RuntimeError(\n",
    "            \"kagglehub not installed and local CSV missing. \"\n",
    "            \"Run `pip install kagglehub[pandas-datasets]` or place the CSV locally.\"\n",
    "        ) from e"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a121a0ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "def clean_text(text):\n",
    "    if pd.isna(text):\n",
    "        return \"\"\n",
    "    return re.sub(r\"\\s+\", \" \", str(text).strip())\n",
    "\n",
    "\n",
    "def parse_list_column(val):\n",
    "    \"\"\"\n",
    "    Columns like ingredients/steps/tags are stored as Python lists (string form).\n",
    "    \"\"\"\n",
    "    if pd.isna(val):\n",
    "        return []\n",
    "    try:\n",
    "        parsed = ast.literal_eval(val)\n",
    "        return [clean_text(x) for x in parsed if isinstance(x, str)]\n",
    "    except Exception:\n",
    "        return []\n",
    "\n",
    "\n",
    "def parse_nutrition(val):\n",
    "    \"\"\"\n",
    "    nutrition is a list of floats in this dataset:\n",
    "    [calories, total fat %, sugar %, sodium %, protein %, sat fat %, carbs %]\n",
    "    \"\"\"\n",
    "    if pd.isna(val):\n",
    "        return {}\n",
    "    try:\n",
    "        parsed = ast.literal_eval(val)\n",
    "        if not isinstance(parsed, (list, tuple)):\n",
    "            return {}\n",
    "        keys = [\n",
    "            \"calories\",\n",
    "            \"total_fat\",\n",
    "            \"sugar\",\n",
    "            \"sodium\",\n",
    "            \"protein\",\n",
    "            \"saturated_fat\",\n",
    "            \"carbohydrates\",\n",
    "        ]\n",
    "        return {k: float(parsed[i]) if i < len(parsed) else None for i, k in enumerate(keys)}\n",
    "    except Exception:\n",
    "        return {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c52bc36f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def structure_recipe(row):\n",
    "    ingredients = parse_list_column(row.get(\"ingredients\"))\n",
    "    steps = parse_list_column(row.get(\"steps\") or row.get(\"directions\"))\n",
    "    tags = parse_list_column(row.get(\"tags\"))\n",
    "    nutrition = parse_nutrition(row.get(\"nutrition\"))\n",
    "    description = clean_text(row.get(\"description\"))\n",
    "    minutes = row.get(\"minutes\")\n",
    "    title = clean_text(row.get(\"name\") or row.get(\"title\") or \"\")\n",
    "\n",
    "    raw_text = (\n",
    "        f\"{title}\"\n",
    "        + \"\\nIngredients:\\n\"\n",
    "        + \"\\n\".join(ingredients)\n",
    "        + \"\\nSteps:\\n\"\n",
    "        + \"\\n\".join(steps)\n",
    "        + \"\\nTags:\\n\"\n",
    "        + \"\\n\".join(tags)\n",
    "    ).strip()\n",
    "\n",
    "    return {\n",
    "        \"id\": str(row.get(\"id\") or row.name),\n",
    "        \"title\": title,\n",
    "        \"description\": description,\n",
    "        \"minutes\": minutes,\n",
    "        \"ingredients\": ingredients,\n",
    "        \"steps\": steps,\n",
    "        \"tags\": tags,\n",
    "        \"nutrition\": nutrition,\n",
    "        \"raw_text\": raw_text,\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5bef777",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Loading raw recipes...\")\n",
    "df = load_raw_recipes()\n",
    "\n",
    "print(\"Structuring & cleaning...\")\n",
    "structured = [structure_recipe(r) for _, r in df.iterrows()]\n",
    "\n",
    "print(\"Selecting top recipes by text length...\")\n",
    "for rec in structured:\n",
    "    rec[\"text_length\"] = len(rec[\"raw_text\"])\n",
    "top_recipes = sorted(structured, key=lambda x: x[\"text_length\"], reverse=True)[:TOP_N]\n",
    "\n",
    "print(f\"Saving {len(top_recipes)} recipes to {OUTPUT_JSON}\")\n",
    "Path(OUTPUT_JSON).parent.mkdir(parents=True, exist_ok=True)\n",
    "with open(OUTPUT_JSON, \"w\", encoding=\"utf-8\") as f:\n",
    "    json.dump(top_recipes, f, indent=2, ensure_ascii=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef228079",
   "metadata": {},
   "source": [
    "## Step 2: Build BM25 Retrieval Engine\n",
    "\n",
    "- Implement BM25-based recipe search using rank_bm25\n",
    "- Support index saving/loading for efficiency\n",
    "- Provide retrieve_bm25(query, corpus, top_n=5) interface\n",
    "- Handle recipe ranking and scoring"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "670448ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_bm25(tokenized_corpus: List[List[str]]) -> BM25Okapi:\n",
    "    return BM25Okapi(tokenized_corpus)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e531434",
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_or_load_bm25(corpus, index_path, text_field=TEXT_FIELD, force_rebuild=False):\n",
    "    if not force_rebuild and os.path.exists(index_path):\n",
    "        with open(index_path, \"rb\") as f:\n",
    "            data = pickle.load(f)\n",
    "        if data.get(\"corpus_size\") == len(corpus):\n",
    "            return data[\"bm25\"], data[\"tok_corpus\"]\n",
    "\n",
    "    tok_corpus = [tokenize(doc.get(text_field, \"\")) for doc in corpus]\n",
    "    bm25 = BM25Okapi(tok_corpus)\n",
    "    with open(index_path, \"wb\") as f:\n",
    "        pickle.dump(\n",
    "            {\"bm25\": bm25, \"tok_corpus\": tok_corpus, \"corpus_size\": len(corpus)},\n",
    "            f,\n",
    "        )\n",
    "    return bm25, tok_corpus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbde56d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Building BM25 index...\")\n",
    "corpus = top_recipes\n",
    "bm25, tok = build_or_load_bm25(corpus, index_path=DEFAULT_INDEX_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c75b845a",
   "metadata": {},
   "source": [
    "#### Span helpers: set kind + conditionally create spans\n",
    "Utility helpers for cleaner tracing:\n",
    "\n",
    "_set_kind(span, kind) safely sets the OpenInference span kind (defaults to CHAIN, falls back to CHAIN if an unknown kind is passed).\n",
    "\n",
    "span_if(enabled, name, *, attrs=None, kind=\"CHAIN\") is a context manager that only creates a span when enabled=True, sets its kind, applies optional attributes, and yields the span. Use it to toggle fine-grained spans without cluttering the code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08b5c924",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _set_kind(span, kind: str = \"CHAIN\"):\n",
    "    span.set_attribute(\n",
    "        SpanAttributes.OPENINFERENCE_SPAN_KIND,\n",
    "        getattr(OpenInferenceSpanKindValues, kind.upper(), OpenInferenceSpanKindValues.CHAIN).value,\n",
    "    )\n",
    "\n",
    "\n",
    "@contextmanager\n",
    "def span_if(enabled: bool, name: str, *, attrs: Dict[str, Any] | None = None, kind: str = \"CHAIN\"):\n",
    "    if not enabled:\n",
    "        yield None\n",
    "        return\n",
    "    with tracer.start_as_current_span(name) as s:\n",
    "        _set_kind(s, kind)\n",
    "        if attrs:\n",
    "            for k, v in attrs.items():\n",
    "                s.set_attribute(k, v)\n",
    "        yield s"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d70f2b25",
   "metadata": {},
   "source": [
    "#### BM25 retriever with tracing\n",
    "@retriever(\"retrieval.bm25\") emits a RETRIEVER span and logs each top-k doc (id, score, short content, metadata) for Phoenix.\n",
    "\n",
    "retrieve_bm25(..., emit_detail_spans=False) runs BM25 and, if enabled, adds tiny child spans (query.normalize → query.tokenize → bm25.score/bm25.rank → results.build).\n",
    "\n",
    "Returns ranked hits with useful metadata (matched terms, doc length, etc.)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "516737ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "def retriever(name: str = \"retrieval.bm25\"):\n",
    "    def outer(fn):\n",
    "        @functools.wraps(fn)\n",
    "        def inner(*a, **kw):\n",
    "            q = kw.get(\"query\", a[0] if a else None)\n",
    "            with tracer.start_as_current_span(name) as s:\n",
    "                s.set_attribute(\n",
    "                    SpanAttributes.OPENINFERENCE_SPAN_KIND,\n",
    "                    OpenInferenceSpanKindValues.RETRIEVER.value,\n",
    "                )\n",
    "                if q is not None:\n",
    "                    s.set_attribute(SpanAttributes.INPUT_VALUE, str(q))\n",
    "                res = fn(*a, **kw)\n",
    "                for i, r in enumerate(res):\n",
    "                    base = f\"retrieval.documents.{i}.document\"\n",
    "                    s.set_attribute(f\"{base}.id\", str(r.get(\"id\")))\n",
    "                    sc = r.get(\"score\")\n",
    "                    if sc is not None:\n",
    "                        s.set_attribute(f\"{base}.score\", float(sc))\n",
    "                    s.set_attribute(\n",
    "                        f\"{base}.content\", (r.get(\"text\", \"\") or \"\")[:1200].replace(\"\\n\", \" \")\n",
    "                    )\n",
    "                    meta = r.get(\"metadata\")\n",
    "                    if meta:\n",
    "                        s.set_attribute(\n",
    "                            f\"{base}.metadata\",\n",
    "                            json.dumps(meta, ensure_ascii=False, separators=(\",\", \":\")),\n",
    "                        )\n",
    "                return res\n",
    "\n",
    "        return inner\n",
    "\n",
    "    return outer\n",
    "\n",
    "\n",
    "@retriever(\"retrieval.bm25\")\n",
    "def retrieve_bm25(\n",
    "    query: str,\n",
    "    corpus: List[Dict[str, Any]],\n",
    "    bm25: BM25Okapi,\n",
    "    tokenized_corpus: List[List[str]],\n",
    "    top_n: int = 5,\n",
    "    text_field: str = TEXT_FIELD,\n",
    "    title_field: str = TITLE_FIELD,\n",
    "    id_field: str = ID_FIELD,\n",
    "    *,\n",
    "    emit_detail_spans: bool = False,\n",
    ") -> List[Dict[str, Any]]:\n",
    "    q_raw = _to_text(query)\n",
    "    with span_if(emit_detail_spans, \"query.normalize\", attrs={\"query.len\": len(q_raw)}):\n",
    "        q_norm = q_raw.strip()\n",
    "\n",
    "    with span_if(emit_detail_spans, \"query.tokenize\"):\n",
    "        qtok = tokenize(q_norm)\n",
    "\n",
    "    with span_if(\n",
    "        emit_detail_spans,\n",
    "        \"bm25.score\",\n",
    "        attrs={\"tokens\": len(qtok), \"N_docs\": len(tokenized_corpus)},\n",
    "    ):\n",
    "        scores = bm25.get_scores(qtok)\n",
    "\n",
    "    with span_if(emit_detail_spans, \"bm25.rank\", attrs={\"k\": int(top_n)}):\n",
    "        top = sorted(range(len(scores)), key=scores.__getitem__, reverse=True)[:top_n]\n",
    "\n",
    "    out: List[Dict[str, Any]] = []\n",
    "    with span_if(emit_detail_spans, \"results.build\", attrs={\"k\": int(top_n)}):\n",
    "        for rank, idx in enumerate(top, 1):\n",
    "            d = corpus[idx]\n",
    "            dt = tokenized_corpus[idx]\n",
    "            out.append(\n",
    "                {\n",
    "                    \"rank\": rank,\n",
    "                    \"score\": float(scores[idx]),\n",
    "                    \"id\": d.get(id_field),\n",
    "                    \"title\": d.get(title_field, \"\"),\n",
    "                    \"text\": d.get(text_field, \"\"),\n",
    "                    \"index\": idx,\n",
    "                    \"metadata\": {\n",
    "                        **{k: _jsonish(v) for k, v in d.items() if k != text_field},\n",
    "                        \"rank\": rank,\n",
    "                        \"bm25_score\": float(scores[idx]),\n",
    "                        \"matched_terms\": sorted(set(qtok) & set(dt)),\n",
    "                        \"doc_len_terms\": len(dt),\n",
    "                    },\n",
    "                }\n",
    "            )\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d002dc59",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"holiday cake pops\"\n",
    "hits = retrieve_bm25(query, corpus, bm25, tok, top_n=5)\n",
    "print(hits)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b40d23c8",
   "metadata": {},
   "source": [
    "## Step 3: Generate Synthetic Queries\n",
    "\n",
    "- Use LLM to generate realistic cooking queries\n",
    "- Focus on complex scenarios requiring specific recipe knowledge\n",
    "- Use ThreadPoolExecutor for parallel processing\n",
    "- Generate 100+ queries with salient facts\n",
    "- Save as data/synthetic_queries.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db294ee7",
   "metadata": {},
   "outputs": [],
   "source": [
    "SYSTEM_PROMPT = \"\"\"You are an advanced user of a recipe search engine.\n",
    "Given a recipe, write ONE realistic, conversational cooking question that depends on a precise, technical detail\n",
    "that is clearly contained in THIS recipe. Focus on:\n",
    "1) Specific methods (e.g., marinate 4 hours, bake at 375°F for 25 minutes)\n",
    "2) Appliance settings (e.g., air fryer 400°F for 12 minutes, pressure cook 8 minutes)\n",
    "3) Ingredient prep details (e.g., slice onions paper-thin, whip cream to soft peaks)\n",
    "4) Timing specifics (e.g., rest dough 30 minutes, simmer 45 minutes)\n",
    "5) Temperature precision (e.g., internal 165°F, oil 350°F)\n",
    "\n",
    "Return EXACTLY a single JSON object:\n",
    "{\"query\":\"...?\",\"salient_fact\":\"<exact quote or tight paraphrase>\"}\"\"\"\n",
    "\n",
    "_TEMPLATE = \"Recipe ID: {id}\\nTitle: {title}\\nKey ingredients: {ingredients}\\nFirst instruction: {first_step}\\n\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e906c86e",
   "metadata": {},
   "source": [
    "#### Parse & normalize LLM output\n",
    "_norm_q(q) — lowercase, strip punctuation, collapse spaces → dedupe-friendly key.\n",
    "\n",
    "_parse_json(text) — parse JSON; if there’s extra text, grab the first {...} block and load it.\n",
    "\n",
    "_parse_and_validate(text) — require query + salient_fact, append “?” if missing, enforce a minimum length, and ensure the fact contains a concrete technical detail (time/temp/ratio/etc.). Returns the cleaned pair or None."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df477810",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _norm_q(q: str) -> str:\n",
    "    s = _to_text(q).lower().translate(_PUNCT).strip()\n",
    "    return re.sub(r\"\\s+\", \" \", s)\n",
    "\n",
    "\n",
    "def _parse_json(text: str) -> Optional[dict]:\n",
    "    s = _to_text(text).strip()\n",
    "    if not s:\n",
    "        return None\n",
    "    try:\n",
    "        return json.loads(s)\n",
    "    except Exception:\n",
    "        m = _JSON_RE.search(s)\n",
    "        return json.loads(m.group(0)) if m else None\n",
    "\n",
    "\n",
    "def _parse_and_validate(text: str) -> Optional[Dict[str, str]]:\n",
    "    obj = _parse_json(text)\n",
    "    if not isinstance(obj, dict):\n",
    "        return None\n",
    "    q = _to_text(obj.get(\"query\")).strip()\n",
    "    f = _to_text(obj.get(\"salient_fact\")).strip()\n",
    "    if not q or not f:\n",
    "        return None\n",
    "    if not q.endswith(\"?\"):\n",
    "        q += \"?\"\n",
    "    if len(q) < 25:\n",
    "        return None\n",
    "    if not _TECH.search(f):\n",
    "        return None\n",
    "    return {\"query\": q, \"salient_fact\": f}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15c02280",
   "metadata": {},
   "source": [
    "Takes a slice of recipes starting at start and grabs batch_rows items (wrapping around if needed).\n",
    "\n",
    "For each recipe, it collects:\n",
    "\n",
    "- id and title\n",
    "\n",
    "- ingredients: the first 6 ingredients, joined as a comma-separated string\n",
    "\n",
    "- first_step: the first instruction, trimmed to 200 characters\n",
    "\n",
    "- _recipe: the full original recipe dict (kept for later)\n",
    "\n",
    "Returns a tidy pandas DataFrame with one row per recipe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33303d18",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _make_batch_df(recipes: List[Dict[str, Any]], start: int, batch_rows: int) -> pd.DataFrame:\n",
    "    rows, n = [], len(recipes)\n",
    "    for i in range(batch_rows):\n",
    "        r = recipes[(start + i) % n]\n",
    "        ings = r.get(\"ingredients\") or []\n",
    "        step0 = (r.get(\"steps\") or [\"\"])[0]\n",
    "        rows.append(\n",
    "            {\n",
    "                \"id\": _to_text(r.get(\"id\")),\n",
    "                \"title\": _to_text(r.get(\"title\")),\n",
    "                \"ingredients\": \", \".join(_to_text(x) for x in ings[:6]),\n",
    "                \"first_step\": _to_text(step0)[:200],\n",
    "                \"_recipe\": r,\n",
    "            }\n",
    "        )\n",
    "    return pd.DataFrame(rows)\n",
    "\n",
    "\n",
    "_PUNCT = str.maketrans(\"\", \"\", string.punctuation)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "985fa787",
   "metadata": {},
   "source": [
    "#### Generate synthetic queries\n",
    "Uses llm_generate to create n unique, validated cooking questions from recipe batches.\n",
    "\n",
    "Dedupes via normalized text, nudges temperature if a batch adds zero.\n",
    "\n",
    "Saves results (with recipe metadata) to out_path and returns the file path."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cffdf24",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_synthetic_queries(\n",
    "    recipes: List[Dict[str, Any]],\n",
    "    n: int = NUM_QUERIES,\n",
    "    out_path: str = SYN_QUERY_PATH,\n",
    "    max_workers: int = MAX_WORKERS,\n",
    "):\n",
    "    Path(out_path).parent.mkdir(parents=True, exist_ok=True)\n",
    "    random.shuffle(recipes)\n",
    "\n",
    "    model = OpenAIModel(\n",
    "        model=QUERY_MODEL,\n",
    "        temperature=0.25,\n",
    "        max_tokens=100,\n",
    "        request_timeout=TIMEOUT_S,\n",
    "    )\n",
    "\n",
    "    unique: Dict[str, Dict[str, Any]] = {}\n",
    "    seen_norm: set[str] = set()\n",
    "    start_idx = 0\n",
    "\n",
    "    with tqdm(total=n, desc=\"synthetic\", mininterval=0.2) as bar:\n",
    "        while len(unique) < n:\n",
    "            need = n - len(unique)\n",
    "            batch_rows = max(256, min(need * 3, 1024))\n",
    "            df = _make_batch_df(recipes, start_idx, batch_rows)\n",
    "            start_idx = (start_idx + batch_rows) % len(recipes)\n",
    "\n",
    "            out = llm_generate(\n",
    "                dataframe=df,\n",
    "                template=_TEMPLATE,\n",
    "                model=model,\n",
    "                system_instruction=SYSTEM_PROMPT,\n",
    "                output_parser=lambda txt, row_idx: (_parse_and_validate(txt) or {}),\n",
    "                run_sync=True,\n",
    "                concurrency=min(max_workers * 3, 48),\n",
    "            )\n",
    "\n",
    "            added = 0\n",
    "            for i, row in out.iterrows():\n",
    "                q = _to_text(row.get(\"query\"))\n",
    "                f = _to_text(row.get(\"salient_fact\"))\n",
    "                if not q or not f:\n",
    "                    continue\n",
    "                key = _norm_q(q)\n",
    "                if key in seen_norm:\n",
    "                    continue\n",
    "\n",
    "                base = df.iloc[i][\"_recipe\"]\n",
    "                unique[q] = {\n",
    "                    \"query\": q,\n",
    "                    \"salient_fact\": f,\n",
    "                    \"source_recipe_id\": base.get(\"id\"),\n",
    "                    \"recipe_name\": _to_text(base.get(\"title\")),\n",
    "                    \"ingredients\": base.get(\"ingredients\", []),\n",
    "                    \"tags\": base.get(\"tags\", []),\n",
    "                }\n",
    "                seen_norm.add(key)\n",
    "                added += 1\n",
    "                if len(unique) >= n:\n",
    "                    break\n",
    "\n",
    "            acc = (added / len(out)) if len(out) else 0.0\n",
    "            bar.set_postfix({\"batch\": len(out), \"added\": added, \"acc%\": f\"{acc * 100:.0f}\"})\n",
    "            if added:\n",
    "                bar.update(added)\n",
    "            else:\n",
    "                model.temperature = min(model.temperature + 0.1, 0.8)\n",
    "\n",
    "    with open(out_path, \"w\", encoding=\"utf-8\") as fp:\n",
    "        json.dump(list(unique.values()), fp, indent=2, ensure_ascii=False)\n",
    "    print(f\"✅ Saved {len(unique)} queries → {out_path}\")\n",
    "    return out_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1beab230",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Generating synthetic queries...\")\n",
    "generate_synthetic_queries(top_recipes)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2774608",
   "metadata": {},
   "source": [
    "#### Upload synthetic queries to Phoenix\n",
    "Loads your JSON of queries, maps to Phoenix-friendly columns: input (query) and output (salient_fact).\n",
    "\n",
    "Flattens optional metadata (recipe_name, source_recipe_id, ingredients, tags) for tracking.\n",
    "\n",
    "Auto-generates a dataset name (or uses dataset_name) and uploads via px.Client().upload_dataset(...).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61cd4acc",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def upload_synthetic_queries_to_phoenix(\n",
    "    json_path: str = SYN_QUERY_PATH,\n",
    "    dataset_name: str | None = None,\n",
    "):\n",
    "    data = json.loads(Path(json_path).read_text(encoding=\"utf-8\"))\n",
    "    df = pd.DataFrame(data)\n",
    "\n",
    "    if \"query\" not in df or \"salient_fact\" not in df:\n",
    "        raise ValueError(\"Expected keys 'query' and 'salient_fact' in the JSON items.\")\n",
    "\n",
    "    df = df.copy()\n",
    "    df[\"input\"] = df[\"query\"].fillna(\"\").astype(str)\n",
    "    df[\"output\"] = df[\"salient_fact\"].fillna(\"\").astype(str)\n",
    "\n",
    "    meta_cols = []\n",
    "    for col in (\"recipe_name\", \"source_recipe_id\"):\n",
    "        if col in df:\n",
    "            df[col] = df[col].fillna(\"\").astype(str)\n",
    "            meta_cols.append(col)\n",
    "    for col in (\"ingredients\", \"tags\"):\n",
    "        if col in df:\n",
    "            df[col] = df[col].apply(\n",
    "                lambda x: json.dumps(x, ensure_ascii=False)\n",
    "                if isinstance(x, (list, dict))\n",
    "                else (\"\" if pd.isna(x) else str(x))\n",
    "            )\n",
    "            meta_cols.append(col)\n",
    "\n",
    "    upload_df = df[[\"input\", \"output\"] + meta_cols]\n",
    "    ds_name = (\n",
    "        dataset_name or f\"recipes-synth-queries-{len(upload_df)}-{time.strftime('%Y%m%d-%H%M%S')}\"\n",
    "    )\n",
    "    px_client = AsyncClient()\n",
    "    ds = await px_client.datasets.create_dataset(\n",
    "        dataframe=upload_df,\n",
    "        name=ds_name,\n",
    "        input_keys=[\"input\"],\n",
    "        output_keys=[\"output\"],\n",
    "        metadata_keys=meta_cols if meta_cols else [],\n",
    "    )\n",
    "    print(f\"✅ Uploaded dataset '{getattr(ds, 'name', ds)}' with {len(upload_df)} rows\")\n",
    "    return ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf68d48d",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Uploading synthetic queries as a dataset in Phoenix...\")\n",
    "await upload_synthetic_queries_to_phoenix(SYN_QUERY_PATH, dataset_name=\"synthetic_queries\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b1c0dce",
   "metadata": {},
   "source": [
    "# Part 2: Evaluate the BM25 Retriever\n",
    "## Step 4: Implement Evaluation\n",
    "\n",
    "- Load synthetic queries and retrieval engine\n",
    "- For each query, run retrieve_bm25() and record results\n",
    "- Calculate standard IR metrics:\n",
    "    - ##Recall@1: Target recipe rank 1\n",
    "    - ##Recall@3: Target recipe in top 3\n",
    "    - ##Recall@5: Target recipe in top 5\n",
    "    - ##MRR: Mean Reciprocal Rank\n",
    "- Save detailed results to results/retrieval_evaluation.json"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24fce350",
   "metadata": {},
   "source": [
    "Let's track our traces in our Phoenix project!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28d813bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_queries_to_phoenix(\n",
    "    items: Union[str, List[Union[str, Dict[str, Any]]]],\n",
    "    corpus,\n",
    "    bm25,\n",
    "    tokenized_corpus,\n",
    "    *,\n",
    "    top_k: int = 5,\n",
    "    tracer=None,\n",
    "):\n",
    "    \"\"\"\n",
    "    For each query, emit a 'rag.query' span; your @retriever-decorated retrieve_bm25()\n",
    "    will emit the child 'retrieval.bm25' span with Phoenix-friendly doc attrs.\n",
    "    \"\"\"\n",
    "    tr = tracer or globals().get(\"tracer\") or register_phoenix()\n",
    "    span_cm = (lambda name: tr.start_as_current_span(name)) if tr else (lambda name: nullcontext())\n",
    "\n",
    "    if isinstance(items, str) and (items.endswith(\".json\") or os.path.exists(items)):\n",
    "        items = json.loads(Path(items).read_text(encoding=\"utf-8\"))\n",
    "\n",
    "    def _txt(x):\n",
    "        return \"\" if x is None else str(x)\n",
    "\n",
    "    def _norm(x):\n",
    "        return \"\" if x is None else str(x)\n",
    "\n",
    "    for it in items:\n",
    "        if isinstance(it, str):\n",
    "            q = it.strip()\n",
    "            salient = None\n",
    "            target_id = None\n",
    "        else:\n",
    "            q = _txt(it.get(\"query\") or it.get(\"input\")).strip()\n",
    "            salient = _txt(it.get(\"salient_fact\") or it.get(\"output\")).strip() or None\n",
    "            target_id = _norm(it.get(\"source_recipe_id\")) or None\n",
    "\n",
    "        with span_cm(\"rag.query\") as root:\n",
    "            if tr:\n",
    "                root.set_attribute(\n",
    "                    SpanAttributes.OPENINFERENCE_SPAN_KIND, OpenInferenceSpanKindValues.CHAIN.value\n",
    "                )\n",
    "                root.set_attribute(SpanAttributes.INPUT_VALUE, q)\n",
    "                root.set_attribute(\"eval.synthetic\", True)\n",
    "                if salient:\n",
    "                    root.set_attribute(\"eval.salient_fact\", salient)\n",
    "                if target_id:\n",
    "                    root.set_attribute(\"eval.target_id\", target_id)\n",
    "                root.set_attribute(\"retriever.top_k\", int(top_k))\n",
    "\n",
    "            t0 = perf_counter()\n",
    "            hits = retrieve_bm25(q, corpus, bm25, tokenized_corpus, top_n=5)\n",
    "            if tr:\n",
    "                root.set_attribute(\"retrieval.latency_ms\", int((perf_counter() - t0) * 1000))\n",
    "                root.set_attribute(\n",
    "                    \"retrieval.top_ids\", json.dumps([_norm(h.get(\"id\")) for h in hits])\n",
    "                )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1dfa8f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Running queries to Phoenix...\")\n",
    "run_queries_to_phoenix(SYN_QUERY_PATH, corpus, bm25, tok, top_k=5, tracer=tracer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5228e181",
   "metadata": {},
   "source": [
    "We can also run our retriever on the synethtic queries dataset we have uploaded."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "771f951c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_bm25_task(\n",
    "    corpus,\n",
    "    bm25,\n",
    "    tokenized_corpus,\n",
    "    k: int,\n",
    "    collector: list,\n",
    "    show_top_k: int = 5,\n",
    "    tracer=None,\n",
    "):\n",
    "    tr = tracer or globals().get(\"tracer\") or register_phoenix()\n",
    "\n",
    "    def _span(name):\n",
    "        return tr.start_as_current_span(name) if tr else nullcontext()\n",
    "\n",
    "    def bm25_task(example):\n",
    "        q = (example[\"input\"].get(\"input\") or example[\"input\"].get(\"query\") or \"\").strip()\n",
    "        salient = (\n",
    "            example[\"output\"].get(\"output\") or example[\"output\"].get(\"salient_fact\") or \"\"\n",
    "        ).strip()\n",
    "        md = example[\"metadata\"] or {}\n",
    "        inp = example[\"input\"] or {}\n",
    "        target_id = _norm_id(md.get(\"source_recipe_id\") or inp.get(\"source_recipe_id\"))\n",
    "\n",
    "        with _span(\"rag.query\") as root:\n",
    "            if tr:\n",
    "                root.set_attribute(\n",
    "                    SpanAttributes.OPENINFERENCE_SPAN_KIND, OpenInferenceSpanKindValues.CHAIN.value\n",
    "                )\n",
    "                root.set_attribute(SpanAttributes.INPUT_VALUE, q)\n",
    "                if salient:\n",
    "                    root.set_attribute(\"eval.salient_fact\", salient)\n",
    "                root.set_attribute(\"retriever.top_k\", int(k))\n",
    "\n",
    "            t0 = perf_counter()\n",
    "            hits = retrieve_bm25(q, corpus, bm25, tokenized_corpus, top_n=5)\n",
    "\n",
    "            keep = min(show_top_k, len(hits))\n",
    "            top = hits[:keep]\n",
    "            top_ids = [_norm_id(h.get(\"id\")) for h in top]\n",
    "            top_titles = [str(h.get(\"title\") or \"\") for h in top]\n",
    "\n",
    "            if tr:\n",
    "                root.set_attribute(\"retrieval.latency_ms\", int((perf_counter() - t0) * 1000))\n",
    "                root.set_attribute(\"retrieval.top_ids\", json.dumps(top_ids))\n",
    "\n",
    "        hit_rank = 0\n",
    "        if target_id:\n",
    "            for r, h in enumerate(hits, 1):\n",
    "                if _norm_id(h.get(\"id\")) == target_id:\n",
    "                    hit_rank = r\n",
    "                    break\n",
    "        r1 = 1.0 if hit_rank == 1 else 0.0\n",
    "        r3 = 1.0 if (hit_rank and hit_rank <= 3) else 0.0\n",
    "        r5 = 1.0 if (hit_rank and hit_rank <= 5) else 0.0\n",
    "        rr = 1.0 / hit_rank if hit_rank else 0.0\n",
    "\n",
    "        collector.append(\n",
    "            {\n",
    "                \"query\": q,\n",
    "                \"salient_fact\": salient,\n",
    "                \"target_id\": target_id or None,\n",
    "                \"top_k\": k,\n",
    "                \"hit_rank\": hit_rank or None,\n",
    "                \"recall@1\": r1,\n",
    "                \"recall@3\": r3,\n",
    "                \"recall@5\": r5,\n",
    "                \"rr\": rr,\n",
    "                \"top5\": [{\"id\": i, \"title\": t} for i, t in zip(top_ids, top_titles)],\n",
    "            }\n",
    "        )\n",
    "\n",
    "        return {\"top_ids\": top_ids, \"top_titles\": top_titles}\n",
    "\n",
    "    return bm25_task"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "951c4635",
   "metadata": {},
   "source": [
    "## 🕰️ Evaluation Time"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3e459d9",
   "metadata": {},
   "source": [
    "#### IR metric evaluators\n",
    "_recall_at_k_generic(k, ...): Checks if the ground-truth recipe ID appears in the model’s top_ids. Returns 1.0 if it’s within the top-k, else 0.0.\n",
    "\n",
    "RecallAt1 / RecallAt3 / RecallAt5: Thin wrappers around the generic function for k = 1, 3, 5.\n",
    "\n",
    "MRR: Mean Reciprocal Rank for a single example. If the ground-truth ID is at rank r in top_ids, returns 1/r; returns 0.0 if it’s not found."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "679803ad",
   "metadata": {},
   "source": [
    "<img alt=\"Document Retrieval Evaluation Image\" src=\"https://storage.googleapis.com/arize-phoenix-assets/assets/images/phoenix-docs-images/ir_metrics_for_rag.png\" width=1000/>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e497d2d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _recall_at_k_generic(\n",
    "    k: int, input=None, output=None, expected=None, metadata=None, reference=None\n",
    ") -> float:\n",
    "    gt = _norm_id((metadata or {}).get(\"source_recipe_id\") or (input or {}).get(\"source_recipe_id\"))\n",
    "    if not gt:\n",
    "        return 0.0\n",
    "    top_ids = (output or {}).get(\"top_ids\", []) if isinstance(output, dict) else []\n",
    "    for rank, did in enumerate(top_ids, 1):\n",
    "        if _norm_id(did) == gt:\n",
    "            return 1.0 if rank <= k else 0.0\n",
    "    return 0.0\n",
    "\n",
    "\n",
    "def RecallAt1(input=None, output=None, expected=None, metadata=None, reference=None) -> float:\n",
    "    return _recall_at_k_generic(\n",
    "        1, input=input, output=output, expected=expected, metadata=metadata, reference=reference\n",
    "    )\n",
    "\n",
    "\n",
    "def RecallAt3(input=None, output=None, expected=None, metadata=None, reference=None) -> float:\n",
    "    return _recall_at_k_generic(\n",
    "        3, input=input, output=output, expected=expected, metadata=metadata, reference=reference\n",
    "    )\n",
    "\n",
    "\n",
    "def RecallAt5(input=None, output=None, expected=None, metadata=None, reference=None) -> float:\n",
    "    return _recall_at_k_generic(\n",
    "        5, input=input, output=output, expected=expected, metadata=metadata, reference=reference\n",
    "    )\n",
    "\n",
    "\n",
    "def MRR(input=None, output=None, expected=None, metadata=None, reference=None) -> float:\n",
    "    gt = _norm_id((metadata or {}).get(\"source_recipe_id\") or (input or {}).get(\"source_recipe_id\"))\n",
    "    if not gt:\n",
    "        return 0.0\n",
    "    top_ids = (output or {}).get(\"top_ids\", []) if isinstance(output, dict) else []\n",
    "    for rank, did in enumerate(top_ids, 1):\n",
    "        if _norm_id(did) == gt:\n",
    "            return 1.0 / rank\n",
    "    return 0.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5eae9830",
   "metadata": {},
   "source": [
    "#### Run BM25 experiment\n",
    "Runs BM25 over a Phoenix dataset by name, logs spans, and scores Recall@1/3/5 + MRR.\n",
    "\n",
    "Collects per-query results, computes a summary, and saves everything to out_path as JSON.\n",
    "\n",
    "Returns the Phoenix experiment and the results file path."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f66460a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def run_bm25_experiment_on_ds(\n",
    "    corpus,\n",
    "    bm25,\n",
    "    tokenized_corpus,\n",
    "    dataset_name: str,\n",
    "    k: int = 50,\n",
    "    show_top_k: int = 5,\n",
    "    out_path: str = EVAL_RESULTS_PATH,\n",
    "):\n",
    "    px_client = AsyncClient()\n",
    "    ds = await px_client.datasets.get_dataset(dataset=dataset_name)\n",
    "\n",
    "    collector = []\n",
    "    task = make_bm25_task(\n",
    "        corpus, bm25, tokenized_corpus, k=k, collector=collector, show_top_k=show_top_k\n",
    "    )\n",
    "\n",
    "    evaluators = [RecallAt1, RecallAt3, RecallAt5, MRR]\n",
    "    experiment = run_experiment(dataset=ds, task=task, evaluators=evaluators)\n",
    "    print(\"✅ Experiment finished.\")\n",
    "\n",
    "    n = len(collector) or 1\n",
    "    sum_r1 = sum(r[\"recall@1\"] for r in collector)\n",
    "    sum_r3 = sum(r[\"recall@3\"] for r in collector)\n",
    "    sum_r5 = sum(r[\"recall@5\"] for r in collector)\n",
    "    sum_rr = sum(r[\"rr\"] for r in collector)\n",
    "\n",
    "    summary = {\n",
    "        \"examples\": n,\n",
    "        \"k\": k,\n",
    "        \"show_top_k\": show_top_k,\n",
    "        \"recall@1\": round(sum_r1 / n, 4),\n",
    "        \"recall@3\": round(sum_r3 / n, 4),\n",
    "        \"recall@5\": round(sum_r5 / n, 4),\n",
    "        \"mrr\": round(sum_rr / n, 4),\n",
    "        \"hit_rate@show_top_k\": round(\n",
    "            sum(1.0 for r in collector if r[\"hit_rank\"] and r[\"hit_rank\"] <= show_top_k) / n, 4\n",
    "        ),\n",
    "    }\n",
    "\n",
    "    payload = {\"summary\": summary, \"results\": collector}\n",
    "    Path(out_path).parent.mkdir(parents=True, exist_ok=True)\n",
    "    Path(out_path).write_text(json.dumps(payload, indent=2, ensure_ascii=False), encoding=\"utf-8\")\n",
    "    print(f\"💾 Saved detailed results → {out_path}\")\n",
    "\n",
    "    return experiment, out_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ac21f91",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Running BM25 experiment on synthetic queries...\")\n",
    "experiment, path = await run_bm25_experiment_on_ds(\n",
    "    corpus, bm25, tok, dataset_name=\"synthetic_queries\", k=TOP_N\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fb280cd",
   "metadata": {},
   "source": [
    "## Optional: Document Retrieval Relevance Evaluator"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d93214f9",
   "metadata": {},
   "source": [
    "<img alt=\"Document Retrieval Evaluation Image\" src=\"https://storage.googleapis.com/arize-phoenix-assets/assets/images/phoenix-docs-images/documentRelevanceDiagram.png\" width=\"1000\"/>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f8592e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "retrieved_documents_df = get_retrieved_documents(\n",
    "    px.Client(), project_name=experiment[\"project_name\"], timeout=None\n",
    ")\n",
    "retrieved_documents_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a43c6df7",
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_model = OpenAIModel(model=\"gpt-4\")\n",
    "relevance_evaluator = RelevanceEvaluator(eval_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2079362b",
   "metadata": {},
   "source": [
    "#### RAG Eval Template =\n",
    "> You are comparing a reference text to a question and trying to determine if the reference text\n",
    ">\n",
    "> contains information relevant to answering the question. Here is the data:\n",
    ">\n",
    ">>    [BEGIN DATA]\n",
    ">>\n",
    ">>    [Question] : query\n",
    ">>\n",
    ">>    [Reference text] : reference\n",
    ">>\n",
    ">>    [END DATA]\n",
    ">\n",
    "> Compare the Question above to the Reference text. You must determine whether the Reference text\n",
    ">\n",
    "> contains information that can answer the Question. Please focus on whether the very specific\n",
    ">\n",
    "> question can be answered by the information in the Reference text.\n",
    ">\n",
    "> Your response must be single word, either \"relevant\" or \"unrelated\",\n",
    ">\n",
    "> and should not contain any text or characters aside from that word.\n",
    ">\n",
    "> \"unrelated\" means that the reference text does not contain an answer to the Question.\n",
    ">\n",
    "> \"relevant\" means the reference text contains an answer to the Question.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "259e930a",
   "metadata": {},
   "outputs": [],
   "source": [
    "retrieved_documents_relevance_df = run_evals(\n",
    "    evaluators=[relevance_evaluator],\n",
    "    dataframe=retrieved_documents_df,\n",
    "    provide_explanation=True,\n",
    "    concurrency=20,\n",
    ")[0]\n",
    "retrieved_documents_relevance_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e64181e",
   "metadata": {},
   "outputs": [],
   "source": [
    "px.Client().log_evaluations(\n",
    "    DocumentEvaluations(eval_name=\"Retrieval Relevance\", dataframe=retrieved_documents_relevance_df)\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
